/*
 * KiHook.cpp
 *
 *  Created on: 2013-10-24
 *      Author: WinDDK
 */

#include "KiHook.h"

ULONG HookedPoint;
ULONG RetPoint;

PSYSTEM_SERVICE_TABLE OldSSDT;
PSYSTEM_SERVICE_TABLE NewSSDT;

PSYSTEM_SERVICE_TABLE OldShadowSSDT;
PSYSTEM_SERVICE_TABLE NewShadowSSDT;

ULONG Filter(ULONG serviceTable, ULONG index, ULONG originalAddress)
{
    char * name = (char*) PsGetCurrentProcess() + 0x174;

    if (!strstr(name, "cheat"))
    {
        return originalAddress;
    }
    if (serviceTable == (ULONG) OldSSDT->ServiceTable)
    {
        if (NewSSDT != NULL)
        {
            return NewSSDT->ServiceTable[index];
        }
    } else
    {
        if (NewShadowSSDT != NULL)
        {
            return NewShadowSSDT->ServiceTable[index];
        }
    }
    return originalAddress;
}

void __declspec(naked) FakeKiFastCallEntry()
{
    _asm
    {
        pushad
        pushfd
        push ebx
        push eax
        push edi
        call Filter
        mov [esp+0x14], eax
        popfd
        popad
        sub esp, ecx
        shr ecx, 2
        jmp RetPoint
    }
}

KiHook::KiHook(ULONG oldSSDT, ULONG newSSDT, ULONG oldShadowSSDT, ULONG newShadowSSDT)
{
    OldSSDT = (PSYSTEM_SERVICE_TABLE)oldSSDT;
    NewSSDT = (PSYSTEM_SERVICE_TABLE)newSSDT;
    OldShadowSSDT = (PSYSTEM_SERVICE_TABLE) oldShadowSSDT;
    NewShadowSSDT = (PSYSTEM_SERVICE_TABLE) newShadowSSDT;

    if (!MmIsAddressValid((PVOID) OldSSDT) || !MmIsAddressValid((PVOID) NewSSDT) || !MmIsAddressValid((PVOID) OldShadowSSDT)
            || !MmIsAddressValid((PVOID) NewShadowSSDT))
    {
        OldSSDT = 0;
        NewSSDT = 0;
        OldShadowSSDT = 0;
        NewShadowSSDT = 0;
    }

    UCHAR signature[] = {0x2B, 0xE1, 0xC1, 0xE9, 0x02};
    UCHAR jumpCode[5];

    ULONG oldFunction = 0;
    ULONG newFunction = 0;
    HookedPoint = 0;


    memcpy(mSignature, signature, sizeof(signature));

    _asm
    {
        mov ecx, 0x176
        rdmsr
        mov oldFunction, eax
    }

    if (!MmIsAddressValid((PVOID) oldFunction))
    {
        return;
    }

    for (int i = 0; i < 0x300; i++)
    {
        PVOID ptr = (PVOID) (oldFunction + i);
        if (!memcmp(ptr, signature, sizeof(signature)))
        {
            HookedPoint = oldFunction + i;
            break;
        }
    }

    if (HookedPoint ==0 )
    {
        return;
    }

    RetPoint = HookedPoint + 5;
    jumpCode[0] = 0xE9;
    *(ULONG*) &jumpCode[1] = (ULONG) FakeKiFastCallEntry - HookedPoint - 5;
    PageProtectOff();
    memcpy((PVOID) HookedPoint, jumpCode, sizeof(jumpCode));
    PageProtectOn();

    KiDebug("Success");
}

KiHook::~KiHook()
{
    if (HookedPoint >0)
    {
        PageProtectOff();
        memcpy((PVOID)HookedPoint, mSignature ,5);
        PageProtectOn();
    }
    KiDebug("Success");
}
